
from __future__ import division, print_function
import os
import glob
import time
import random
import argparse
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from datetime import datetime
from dateutil.relativedelta import relativedelta
import warnings
warnings.filterwarnings("ignore")

from utils import *
from models import GAT, ResGAT, MLP, EnhancedMLP

data_path = fr'{os.getcwd()}/data'
model_path = fr'{os.getcwd()}/model/model_saved'

# 目前仅接受x->y和x->rank(y)
class Args:
    no_cuda = False
    res = False
    seed = 42
    epochs = 200
    lr = 1e-3
    weight_decay = 5e-4
    hidden = 128 
    nb_heads = 2 
    dropout_input = 0.5
    dropout_hidden = 0.2
    alpha = 0.2
    patience = 20
    loss = 'huber'
    delta = 1.0
    model = 'nn'
    dataset = 'stock_node_vwap_edge_sec'
    use_all_ones_adj = False

args = Args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
print('是否使用GPU: ',args.cuda)

random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)


# 修改load_monthly_data函数以支持选择性加载adj
def load_monthly_data(start_month, end_month, load_adj=True):
    """修改后的数据加载函数，支持选择性加载邻接矩阵"""
    idx_list, features_list, adjs_list, labels_list = [], [], [], []
    current = datetime.strptime(start_month, "%Y%m")
    end = datetime.strptime(end_month, "%Y%m")
    
    while current <= end:
        month_str = current.strftime("%Y%m")
        idx, adj, features, labels = load_data(path=data_path, 
                            dataset=args.dataset, month=month_str,
                            use_all_ones_adj = args.use_all_ones_adj)
        idx_list.append(idx)
        features_list.append(features)
        if load_adj:
            adjs_list.append(adj)
        labels_list.append(labels)
        current = datetime(current.year + (current.month // 12), (current.month % 12) + 1, 1)
    
    if load_adj:
        return idx_list, features_list, adjs_list, labels_list
    return idx_list, features_list, labels_list


def shift_month(date, months):
    year = date.year + (date.month + months - 1) // 12
    month = (date.month + months - 1) % 12 + 1
    return datetime(year, month, 1)


def train_generic_model(model_type):
    """通用训练函数，支持GAT和MLP模型"""
    start_test = datetime.strptime("201801", "%Y%m")
    end_test = datetime.strptime("202412", "%Y%m")
    cur_test = start_test
    round_idx = 0
    all_time_test_loss = []
    all_time_test_IC = []

    while cur_test <= end_test:
        # 准备数据周期
        train_start = shift_month(cur_test, -72)
        train_end = shift_month(cur_test, -13)
        val_start = shift_month(cur_test, -12)
        val_end = shift_month(cur_test, -1)
        test_start = cur_test
        test_end = shift_month(cur_test, 11)

        print(f"\n--- Round {round_idx} ---")
        print(f"Train: {train_start:%Y%m}-{train_end:%Y%m}, Val: {val_start:%Y%m}-{val_end:%Y%m}, Test: {test_start:%Y%m}-{test_end:%Y%m}")

        # 加载数据（根据模型类型决定是否加载adj）
        load_adj = model_type == 'gat'
        train_data = load_monthly_data(train_start.strftime("%Y%m"), train_end.strftime("%Y%m"), load_adj)
        val_data = load_monthly_data(val_start.strftime("%Y%m"), val_end.strftime("%Y%m"), load_adj)
        test_data = load_monthly_data(test_start.strftime("%Y%m"), test_end.strftime("%Y%m"), load_adj)
        
        if load_adj:
            train_idx, train_features, train_adjs, train_labels = train_data
            val_idx, val_features, val_adjs, val_labels = val_data
            test_idx, test_features, test_adjs, test_labels = test_data
        else:
            train_idx, train_features, train_labels = train_data
            val_idx, val_features, val_labels = val_data
            test_idx, test_features, test_labels = test_data
            train_adjs = val_adjs = test_adjs = [None] * len(train_features)  # 填充None保持结构一致

        # 转移到CUDA
        device = torch.device('cuda' if args.cuda else 'cpu')
            
        # 转移到设备(CUDA或CPU)
        def to_device(*tensors):
            return [t.to(device) for t in tensors]
            
        train_features = to_device(*train_features)
        train_labels = to_device(*train_labels)
        val_features = to_device(*val_features)
        val_labels = to_device(*val_labels)
        test_features = to_device(*test_features)
        test_labels = to_device(*test_labels)

        if load_adj:
            train_adjs = to_device(*train_adjs)
            val_adjs = to_device(*val_adjs)
            test_adjs = to_device(*test_adjs)

        # 转换为Variable
        def to_variable(*tensors):
            return [Variable(t.float()) for t in tensors]
            
        train_features = to_variable(*train_features)
        train_labels = to_variable(*train_labels)
        val_features = to_variable(*val_features)
        val_labels = to_variable(*val_labels)
        test_features = to_variable(*test_features)
        test_labels = to_variable(*test_labels)

        if load_adj:
            train_adjs = to_variable(*train_adjs)
            val_adjs = to_variable(*val_adjs)
            test_adjs = to_variable(*test_adjs)
            
        # 初始化模型
        if model_type == 'gat':
            if args.res:
                model = ResGAT(
                    nfeat=train_features[0].shape[1], 
                    nhid=args.hidden, 
                    nclass=1,
                    dropout_input=args.dropout_input,
                    dropout_hidden=args.dropout_hidden,
                    nheads=args.nb_heads, 
                    alpha=args.alpha
                )
            else:
                model = GAT(
                    nfeat=train_features[0].shape[1], 
                    nhid=args.hidden, 
                    nclass=1,
                    dropout_input=args.dropout_input,
                    dropout_hidden=args.dropout_hidden,
                    nheads=args.nb_heads, 
                    alpha=args.alpha
                )
        else:  # model_type == 'nn'
            # model = EnhancedMLP(
            #     nfeat=train_features[0].shape[1]
            # )
            model = MLP(
                nfeat=train_features[0].shape[1],
                nhid=args.hidden,
                dropout_input=args.dropout_input,
                dropout_hidden=args.dropout_hidden,
                nheads=args.nb_heads
            )
        
        model.to(device)

        optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

        # 训练循环
        best_epoch, loss_values, best = 0, [], args.epochs + 1
        bad_counter = 0
        train_loss_curve, val_loss_curve = [], []
        train_IC_curve, val_IC_curve = [], []

        for epoch in range(args.epochs):
            epoch_start_time = time.time()

            # 训练阶段
            model.train()
            optimizer.zero_grad()
            losses, ICs = [], []
            for m in range(len(train_features)):
                # 根据模型类型调用不同的forward方法
                if model_type == 'gat':
                    output = model(train_features[m], train_adjs[m])
                else:
                    output = model(train_features[m])
                
                if args.loss == 'huber':
                    loss = huber_loss(output, train_labels[m], args.delta)
                else:
                    loss = F.mse_loss(output, train_labels[m])
                
                if 'rank' in args.dataset:
                    IC = pearson_IC(rank(output), train_labels[m])
                else:
                    IC = spearman_IC(output, train_labels[m])
                
                losses.append(loss)
                ICs.append(IC)
            
            avg_loss = torch.stack(losses).mean()
            avg_loss.backward()
            optimizer.step()

            # 验证阶段
            val_losses, val_ICs = [], []
            model.eval()
            for m in range(len(val_features)):
                if model_type == 'gat':
                    output = model(val_features[m], val_adjs[m])
                else:
                    output = model(val_features[m])
                
                if args.loss == 'huber':
                    loss = huber_loss(output, val_labels[m], args.delta)
                else:
                    loss = F.mse_loss(output, val_labels[m])
                
                if "rank" in args.dataset:
                    IC = pearson_IC(rank(output), val_labels[m])
                else:
                    IC = spearman_IC(output, val_labels[m])
                
                val_losses.append(loss)
                val_ICs.append(IC)

            val_loss = torch.stack(val_losses).mean().item()
            avg_IC_train = torch.stack(ICs).mean().item()
            avg_IC_val = torch.stack(val_ICs).mean().item()

            # 记录曲线
            train_loss_curve.append(avg_loss.item())
            val_loss_curve.append(val_loss)
            train_IC_curve.append(avg_IC_train)
            val_IC_curve.append(avg_IC_val)

            # 打印进度
            epoch_duration = time.time() - epoch_start_time
            print(f"Epoch: {epoch + 1:04d} loss_train: {avg_loss:.4f} IC_train: {avg_IC_train:.4f} "
                  f"loss_val: {val_loss:.4f} IC_val: {avg_IC_val:.4f} time: {epoch_duration:.4f}s")

            # 保存模型和早停判断
            model_dir = 'gat' if model_type == 'gat' else 'nn'
            torch.save(model.state_dict(), f'{model_path}/{model_dir}/model_{round_idx}_{epoch}.pt')
            
            if val_loss < best:
                best, best_epoch, bad_counter = val_loss, epoch, 0
            else:
                bad_counter += 1
                if bad_counter >= args.patience:
                    break
        
        make_loss_figure(train_loss_curve, val_loss_curve, round_idx, args.model, args.res)
        make_IC_figure(train_IC_curve, val_IC_curve, round_idx, args.model, args.res)

        # 测试阶段
        model_dir = 'gat' if model_type == 'gat' else 'nn'
        model.load_state_dict(torch.load(f'{model_path}/{model_dir}/model_{round_idx}_{best_epoch}.pt'))
        model.eval()

        # 保存结果
        out_dir = f'results/{args.model}/Res_{args.res}/results_round_{round_idx}' if model_type == 'gat' else f'results/{args.model}/results_round_{round_idx}'
        os.makedirs(out_dir, exist_ok=True)
        
        outputs, losses, ICs = [], [], []
        cur_date = cur_test
        for m in range(len(test_features)):
            if model_type == 'gat':
                output = model(test_features[m], test_adjs[m])
            else:
                output = model(test_features[m])
            
            if args.loss == 'huber':
                loss = huber_loss(output, test_labels[m], args.delta)
            else:
                loss = F.mse_loss(output, test_labels[m])
            if "rank" in args.dataset:
                IC = pearson_IC(rank(output), test_labels[m])
            else:
                IC = spearman_IC(output, test_labels[m])
            
            outputs.append(output.detach().cpu().numpy())
            losses.append(loss.item())
            ICs.append(IC.item())

            # 保存每月结果
            df = pd.DataFrame({
                'index': test_idx[m],
                'prediction': output.detach().cpu().numpy().squeeze()
            })
            df.to_csv(f'{out_dir}/{cur_date.strftime("%Y-%m")}.csv', index=False)

            # 推进一个月
            cur_date += relativedelta(months=1)
        
        torch.save(model.state_dict(), os.path.join(out_dir, 'best_model.pt'))
        
        # 记录结果
        avg_test_loss = np.mean(losses)
        avg_test_IC = np.mean(ICs)
        all_time_test_loss.append(avg_test_loss)
        all_time_test_IC.append(avg_test_IC)
        
        print(f"Round {round_idx} complete. Avg Test Loss: {avg_test_loss:.4f}, Avg IC: {avg_test_IC:.4f}")
        
        # 下一轮
        cur_test = shift_month(cur_test, 12)
        round_idx += 1
    
    print('All rounds test loss:', all_time_test_loss)
    print('All rounds test IC:', all_time_test_IC)


# 主执行部分
if __name__ == '__main__':
    if args.model in ['gat', 'nn']:
        if args.model == 'gat':
            print(f"ResNet: {args.res}")
        print(f'Training {args.model.upper()} model...')
        train_generic_model(args.model)
    else:
        raise ValueError(f"Invalid model type: '{args.model}'. Expected 'gat' or 'nn'.")
